import pytesseract
from PIL import Image
import cv2
import numpy as np

def find_text_by_ocr(image_path: str, target_text: str = "基础数据") -> bool:
    """
    使用 OCR 识别截图中是否包含指定文字
    :param image_path: 截图路径
    :param target_text: 要查找的目标文字
    :return: 是否找到目标文字
    """
    # 读取图像
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"无法加载图像: {image_path}")

    # 转为 PIL 图像格式
    pil_img = Image.fromarray(img)
    pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
    
    # 使用 OCR 识别文字（支持中文）并获取位置信息
    try:
        # 尝试不同的页面分割模式以更好地识别连续文本
        data = pytesseract.image_to_data(pil_img, lang='chi_sim+eng', output_type=pytesseract.Output.DICT, config='--psm 6')
        n_boxes = len(data['level'])
        
        # 合并同一行上的文本以更好地识别连续文本
        lines = merge_text_lines(data)
        print("合并后的行文本:")
        for i, line in enumerate(lines):
            # 移除合并后行文本中的所有空格，使文本紧密连接
            cleaned_text = ''.join(line['text'].split())
            print(f"行 {i+1}: '{cleaned_text}' 位置: (x={line['x']}, y={line['y']}, width={line['w']}, height={line['h']})")
        
        # 查找目标文本（在合并行中查找）
        found = False
        found_position = None
        # 在合并后的行中查找目标文本
        for i, line in enumerate(lines):
            # 移除文本中的所有空格后再查找
            cleaned_text = ''.join(line['text'].split())
            # 只要包含目标文本就匹配
            if target_text in cleaned_text:
                print(f"在合并行中找到目标文字: '{target_text}' 位置: (x={line['x']}, y={line['y']}, width={line['w']}, height={line['h']})")
                found = True
                found_position = (int(line['x']), int(line['y']), int(line['w']), int(line['h']))
                break  # 找到第一个匹配的就停止
                
        if not found:
            print(f"❌ 未找到目标文字: '{target_text}'")
            return False
            
        # 在图像上绘制找到的文本框
        if found_position:
            x, y, w, h = found_position
            # 在图像上绘制矩形框（绿色，线宽2）
            cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)
            # 在框上方添加文本标签
            cv2.putText(img, target_text, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
            
            # 保存带标注的图像
            output_path = image_path.replace('.png', '_marked.png').replace('.jpg', '_marked.jpg')
            cv2.imwrite(output_path, img)
            print(f"已保存带标注的图像到: {output_path}")
            
        return True
    except Exception as e:
        print(f"OCR 识别失败: {e}")
        return False

def merge_text_lines(data, y_tolerance=10):
    """
    合并同一行上的文本框
    
    :param data: pytesseract 返回的数据字典
    :param y_tolerance: Y坐标容差，用于判断是否在同一行
    :return: 合并后的行列表
    """
    lines = []
    n_boxes = len(data['level'])
    
    # 提取有效的文本框
    boxes = []
    for i in range(n_boxes):
        text = data['text'][i].strip()
        # 去除空字符串但仍显示空白字符
        if len(text) > 0:
            box = {
                'text': text,
                'x': data['left'][i],
                'y': data['top'][i],
                'w': data['width'][i],
                'h': data['height'][i]
            }
            boxes.append(box)
    
    # 根据Y坐标对文本框进行分组
    if not boxes:
        return lines
        
    # 按Y坐标排序
    boxes.sort(key=lambda b: b['y'])
    
    # 分组同一行的文本框
    current_line = [boxes[0]]
    for i in range(1, len(boxes)):
        # 如果当前文本框与前一个文本框在同一行，则添加到当前行
        if abs(boxes[i]['y'] - boxes[i-1]['y']) <= y_tolerance:
            current_line.append(boxes[i])
        else:
            # 否则，结束当前行并将文本框按X坐标排序后合并
            current_line.sort(key=lambda b: b['x'])
            merged_text = ' '.join([b['text'] for b in current_line])
            
            # 计算合并后的边界框
            min_x = min([b['x'] for b in current_line])
            max_x = max([b['x'] + b['w'] for b in current_line])
            min_y = min([b['y'] for b in current_line])
            max_y = max([b['y'] + b['h'] for b in current_line])
            
            lines.append({
                'text': merged_text,
                'x': min_x,
                'y': min_y,
                'w': max_x - min_x,
                'h': max_y - min_y
            })
            
            # 开始新行
            current_line = [boxes[i]]
    
    # 处理最后一行
    if current_line:
        current_line.sort(key=lambda b: b['x'])
        merged_text = ' '.join([b['text'] for b in current_line])
        
        # 计算合并后的边界框
        min_x = min([b['x'] for b in current_line])
        max_x = max([b['x'] + b['w'] for b in current_line])
        min_y = min([b['y'] for b in current_line])
        max_y = max([b['y'] + b['h'] for b in current_line])
        
        lines.append({
            'text': merged_text,
            'x': min_x,
            'y': min_y,
            'w': max_x - min_x,
            'h': max_y - min_y
        })
    
    return lines

# 示例调用
if __name__ == '__main__':
    screenshot_path = r'D:\PycharmProjects\auto-script\screentshot\APP-Open.png'
    found = find_text_by_ocr(screenshot_path, "内打开")
    if found:
        print("可以继续点击操作")
    else:
        print("等待页面加载完成...")